import cv2
import numpy as np
from volleyball_detection import VolleyballDetector

def draw_detection_zone(frame, zone, color=(0, 255, 0), thickness=2):
    """
    在视频帧上绘制检测区域
    :param frame: 视频帧
    :param zone: 检测区域坐标
    :param color: 绘制颜色
    :param thickness: 线条粗细
    """
    # 绘制矩形框
    cv2.rectangle(frame, 
                 (zone['left'], zone['top']), 
                 (zone['right'], zone['bottom']), 
                 color, 
                 thickness)
    
    # 添加半透明遮罩
    overlay = frame.copy()
    cv2.rectangle(overlay, 
                 (zone['left'], zone['top']), 
                 (zone['right'], zone['bottom']), 
                 color, 
                 -1)
    cv2.addWeighted(overlay, 0.1, frame, 0.9, 0, frame)

def draw_detection_results(frame, results, color=(0, 255, 0), thickness=2):
    """
    在视频帧上绘制检测结果
    :param frame: 视频帧
    :param results: 检测结果
    :param color: 绘制颜色
    :param thickness: 线条粗细
    """
    # 绘制人物框
    if results['person']:
        person_bbox = results['person']['bbox']
        cv2.rectangle(frame, 
                     (person_bbox[0], person_bbox[1]), 
                     (person_bbox[2], person_bbox[3]), 
                     (255, 0, 0), 
                     thickness)
        cv2.putText(frame, 
                   f"Person: {results['person']['confidence']:.2f}", 
                   (person_bbox[0], person_bbox[1] - 10),
                   cv2.FONT_HERSHEY_SIMPLEX, 
                   0.5, 
                   (255, 0, 0), 
                   2)

    # 绘制排球框
    if results['volleyball']:
        volleyball_bbox = results['volleyball']['bbox']
        cv2.rectangle(frame, 
                     (volleyball_bbox[0], volleyball_bbox[1]), 
                     (volleyball_bbox[2], volleyball_bbox[3]), 
                     (0, 0, 255), 
                     thickness)
        cv2.putText(frame, 
                   f"Volleyball: {results['volleyball']['confidence']:.2f}", 
                   (volleyball_bbox[0], volleyball_bbox[1] - 10),
                   cv2.FONT_HERSHEY_SIMPLEX, 
                   0.5, 
                   (0, 0, 255), 
                   2)

    # 显示计数
    cv2.putText(frame, 
               f"Count: {results['count']}", 
               (10, 30),
               cv2.FONT_HERSHEY_SIMPLEX, 
               1, 
               (0, 255, 0), 
               2)

def display_video_stream(source=0):
    """
    实时显示视频流和检测结果（不保存）
    :param source: 视频源（可以是摄像头索引或视频文件路径）
    """
    # 创建检测器实例
    detector = VolleyballDetector(0.1, 0.2, 0.3, 0.8)
    
    # 打开视频源
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频源: {source}")
    
    # 获取视频的FPS
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_delay = int(1000 / fps)  # 毫秒
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 进行检测
            results = detector.detect_objects(frame)
            
            # 绘制检测区域
            draw_detection_zone(frame, results['detection_zone'])
            
            # 绘制检测结果
            draw_detection_results(frame, results)
            
            # 显示视频帧
            cv2.imshow('Volleyball Detection (Stream Mode)', frame)
            
            # 按'q'退出
            if cv2.waitKey(frame_delay) & 0xFF == ord('q'):
                break
                
    finally:
        # 释放资源
        cap.release()
        cv2.destroyAllWindows()

def display_video_save(source=0, output_filename="output.mp4"):
    """
    显示视频流和检测结果，并保存到文件
    :param source: 视频源（可以是摄像头索引或视频文件路径）
    :param output_filename: 输出视频文件名
    """
    # 创建检测器实例
    detector = VolleyballDetector(0.1, 0.2, 0.3, 0.8)
    
    # 打开视频源
    cap = cv2.VideoCapture(source)
    if not cap.isOpened():
        raise ValueError(f"无法打开视频源: {source}")
    
    # 获取视频的FPS和尺寸
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frame_delay = int(1000 / fps)  # 毫秒
    
    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_filename, fourcc, fps, (frame_width, frame_height))
    print(f"视频将保存为: {output_filename}")
    
    try:
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            # 进行检测
            results = detector.detect_objects(frame)
            
            # 绘制检测区域
            draw_detection_zone(frame, results['detection_zone'])
            
            # 绘制检测结果
            draw_detection_results(frame, results)
            
            # 保存处理后的帧
            out.write(frame)
            
            # 显示视频帧
            cv2.imshow('Volleyball Detection (Save Mode)', frame)
            
            # 按'q'退出
            if cv2.waitKey(frame_delay) & 0xFF == ord('q'):
                break
                
    finally:
        # 释放资源
        cap.release()
        out.release()
        cv2.destroyAllWindows()

def main():
    # 可以选择使用摄像头（0表示默认摄像头）或视频文件
    source = 0  # 使用摄像头
    source = "1.mp4"  # 使用视频文件
    
    # 选择模式：stream 或 save
    mode = "stream"  # 或 "save"
    
    try:
        if mode == "stream":
            print("启动流式显示模式（不保存视频）...")
            display_video_stream(source)
        else:
            print("启动保存模式...")
            display_video_save(source, "output.mp4")
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == "__main__":
    main() 